package com.hanliukui.ai.langchain4j;


import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;
import org.springframework.beans.BeansException;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;

import java.util.List;

@SpringBootApplication
public class SmartAiApplication implements CommandLineRunner, ApplicationContextAware {

    private ApplicationContext applicationContext;

    public static void main(String[] args) {
        SpringApplication.run(SmartAiApplication.class, args);

    }

    @Override
    public void run(String... args) throws Exception {
        EmbeddingModel embeddingModel = applicationContext.getBean(EmbeddingModel.class);
        EmbeddingStore embeddingStore = applicationContext.getBean(EmbeddingStore.class);

        // 这里在初始的时候，初始化一些向量数据
        EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(DocumentSplitters.recursive(300, 30))
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        initEmbeddingKnowledge(embeddingStoreIngestor);
    }

    private static void initEmbeddingKnowledge(EmbeddingStoreIngestor embeddingStoreIngestor) {
        // 初始化一些向量数据
        List<Document> documents = ClassPathDocumentLoader.loadDocuments("knowledges", new TextDocumentParser());

        IngestionResult ingest = embeddingStoreIngestor
                .ingest(documents);
        System.out.println(ingest.tokenUsage());
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }
}